# -*- coding: utf-8 -*-

"""给定一个链表数组，每个链表都已经按升序排列。
请将所有链表合并到一个升序链表中，返回合并后的链表。

示例 1：
输入：lists = [[1,4,5],[1,3,4],[2,6]]
输出：[1,1,2,3,4,4,5,6]
解释：链表数组如下：
[
  1->4->5,
  1->3->4,
  2->6
]
将它们合并到一个有序链表中得到。
1->1->2->3->4->4->5->6

示例 2：
输入：lists = []
输出：[]

示例 3：
输入：lists = [[]]
输出：[]

提示：
k == lists.length
0 <= k <= 10^4
0 <= lists[i].length <= 500
-10^4 <= lists[i][j] <= 10^4
lists[i] 按 升序 排列
lists[i].length 的总和不超过 10^4"""

from sortedcontainers import SortedList

# Definition for singly-linked list.
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

class Solution:
    """将每个链表的头节点以 (val, idx) := (链表头节点值, 链表在数组中的位置) 的形式放到有序集合中
    从有序集合中删除返回最小元素，
    链接到answer
    所在链表删除首节点
    如果所在链表还没有到达尾部，将该链表首节点 (val, idx) 加入有序集合中
    ...
    循环操作，直至有序集合为空
    
    实际来看，用小顶堆比有序集合效果更好"""
    def mergeKLists(self, lists) -> ListNode:
        answer = ListNode()
        
        rbtree = SortedList()
        for i, node in enumerate(lists):
            if node:
                rbtree.add((node.val, i))

        p = answer
        while rbtree:
            rbnode = rbtree[0]
            del rbtree[0]
            
            p.next = ListNode(rbnode[0])
            p = p.next

            idx = rbnode[1]
            lists[idx] = lists[idx].next

            if lists[idx]:
                rbtree.add((lists[idx].val, idx))

        return answer.next


if __name__ == '__main__':
    so = Solution()
    answer = so.mergeKLists([ListNode(1, ListNode(4, ListNode(5))), ListNode(1, ListNode(3, ListNode(4))), ListNode(2, ListNode(6))])
    while answer:
        print(answer.val, sep=', ')
        answer=answer.next
